FROM quay.io/ascend/cann:8.2.rc1.alpha002-910b-openeuler22.03-py3.10

# 设置环境变量
ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=1 \
    PYTHONDONTWRITEBYTECODE=1 \
    ASCEND_TOOLKIT_HOME=/usr/local/Ascend/ascend-toolkit/latest \
    LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64:/usr/local/Ascend/driver/lib64/stub:/usr/local/Ascend/ascend-toolkit/latest/lib64:/usr/local/Ascend/add-ons:$LD_LIBRARY_PATH

# 安装系统依赖
RUN yum update -y && \
    yum install -y \
    git \
    wget \
    vim \
    gcc \
    gcc-c++ \
    make \
    cmake \
    unzip \
    libglvnd-glx \
    mesa-libGL \
    libX11 \
    libXext \
    libXrender \
    && yum clean all

WORKDIR /app

# 设置 Python 镜像源
RUN pip config set global.index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple

# 安装 PyTorch 和 torch_npu (适配 CANN 8.0+)
# 注意：torch-npu 版本需要与 torch 版本匹配
# 
# 选项 1: torch-npu 2.6.0 (最新版本，可能不稳定)
# - torch==2.6.0
# - torchvision==0.21.0 (torch 2.6.0 对应的版本)
# - torchaudio==2.6.0
# RUN pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0
# RUN pip install torch-npu==2.6.0
# 
# 选项 2: torch-npu 2.3.1 (稳定版本，推荐使用)
# torch 2.1.0 与 transformers 不兼容（_pytree API 问题）
# 使用 2.3.1 版本可以解决兼容性问题
RUN pip install torch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1
RUN pip install torch-npu==2.3.1

# 安装项目依赖
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# 确保 numpy 版本兼容 torch（其他包可能升级 numpy 到 2.x）
RUN pip install --no-cache-dir "numpy<2.0.0"

# 修复 diffusers 的 NPU 兼容性问题
COPY patch_diffusers_npu.py /app/patch_diffusers_npu.py
COPY patch_attention_npu.py /app/patch_attention_npu.py
COPY check_patch.py /app/check_patch.py
RUN python3 /app/patch_diffusers_npu.py
RUN python3 /app/patch_attention_npu.py
RUN python3 /app/check_patch.py
RUN python3 -c "import diffusers, os; f=os.path.join(os.path.dirname(diffusers.__file__), 'models/transformers/transformer_z_image.py'); content=open(f).read(); assert '_npu_patched' in content or '_freqs_cis_npu_converted' in content, 'Patch verification failed'; assert 'result.append(self.freqs_cis[i][index])' not in content, 'Original line still exists!'; print('Patch verified successfully')"

# 复制项目代码
COPY api/ /app/api/
COPY config/ /app/config/

# 设置环境变量
ENV PYTHONPATH=/app:$PYTHONPATH
ENV ASCEND_OPP_PATH=/usr/local/Ascend/ascend-toolkit/latest/opp

# 暴露端口
EXPOSE 8000

# 启动命令
# 使用 uvicorn 模块方式运行，支持相对导入
CMD ["python3", "-m", "uvicorn", "api.main:app", "--host", "0.0.0.0", "--port", "8000"]

